from subway_qa.agents.extraction import extract
from subway_qa.db import search
from subway_qa.gaode import AMapSubwayPlanner

planner = AMapSubwayPlanner()


PROMPT = """你是地铁智能客服，你需要根据检索到的知识回答用户的问题
<knowledge>
{knowledge}
</knowledge>
对于无关的问题或者不在检索到的知识范围内的问题你可以拒答。
"""


def build_prompt(query, use_query=False, extraction_cache={}, history=[]):
    extraction = extract(history)
    extraction_cache["is_relevant"] = extraction["is_relevant"]
    if extraction["start"]:
        extraction_cache["start"] = extraction["start"]
    if extraction["end"]:
        extraction_cache["end"] = extraction["end"]

    if extraction_cache["is_relevant"]:
        if "start" not in extraction_cache or not extraction_cache["start"]:
            return "请你扮演一个地铁智能客服，询问我出发站点是哪里"
        if "end" not in extraction_cache or not extraction_cache["end"]:
            return "请你扮演一个地铁智能客服，询问我目的地是哪里"
        start_loc = planner.get_coordinates(extraction_cache["start"] + "地铁站")
        end_loc = planner.get_coordinates(extraction_cache["end"] + "地铁站")
        route_data = planner.get_subway_route(start_loc, end_loc)
        summary = planner.get_route_summary(route_data)
        return "请用语音播报下面的路线规划, 不要说其他额外的东西：" + summary
    else:
        qa_pair = search(query)

        if use_query:
            return PROMPT.format(query=query, knowledge=qa_pair) + f"\n\nuser: {query}"
        else:
            return PROMPT.format(knowledge=qa_pair)